import os

import cv2
from flask_restx import Api
import numpy as np

import json
import platform
import shutil
from get_video_code import convert_video_to_h264
from result_true import process_annotation_files
from flask_cors import *
from flask import Flask, request, jsonify
import os.path
from gevent import pywsgi
import yaml
from ultralytics import YOLO
from heavy import fake_heavy_task
model_link = r'best20241230.pt'#通用模型
model_link1= r'fish20241229.pt' #钓鱼模型
pixel_to_cm = 4.2
sys = platform.system()


def load_yaml(file_path):
    """加载yaml文件并返回其中的 'names' 内容"""
    with open(file_path, 'r', encoding='utf-8') as f:
        result_dict = yaml.load(f.read(), Loader=yaml.FullLoader)
        return result_dict['names']
app = Flask(__name__)
api = Api(app)

def load_model(model_type=0):  # 默认model_type=0
    # 根据模型类型选择对应的模型
    if model_type == 1:
        model = YOLO(model_link1)  # 加载钓鱼模型
        model_config = 'fish.yaml'
    else:
        model = YOLO(model_link)  # road.yaml需要改名
        model_config = 'road.yaml'

    # 加载配置文件
    result_dict = load_yaml(model_config)

    # 返回模型和配置
    return model, result_dict
def predict_img(image_path,model_type):
    #model = YOLO(model_link)


    model,result_dict= load_model(model_type)  # 调用时传入 model_type
    if sys == "Windows":
        dir_path = '\\'.join(image_path.split('\\')[:-1])
        image_name = image_path.split('\\')[-1]
    else:
        dir_path = '/'.join(image_path.split('/')[:-1])
        image_name = image_path.split('/')[-1]
    # results = model.predict(device="cpu", source=image_path, save=True,show_labels=True,show_conf=True,imgsz=640,line_width=10,save_txt=True)
    results = model.predict(image_path, device='cpu', save=True, save_txt=True, conf=0.3,  iou=0.2)

    save_dir = results[0].save_dir

    image_folder = save_dir
    annotation_folder = os.path.join(save_dir, 'labels')
    process_annotation_files(image_folder, annotation_folder)
    result_txt = os.path.join(save_dir, 'labels', f"{image_name.split('.')[0]}_transformed.txt")
    data = {
        "file_path": dir_path,
        "correct_results": [
        ],
    }
    try:
        with open(result_txt) as fp:
            predict_result = fp.readlines()
        current_path = os.path.dirname(__file__)
        image_link = os.path.abspath(os.path.join(current_path, image_name))
        predict_result = [result.replace('\n', '').split(' ') for result in predict_result]
    except:
        predict_result = []
    for result in predict_result:
        bbox = {
            "[x1][y1]": [float(result[1]), float(result[2])],
            "[x2][y2]": [float(result[3]), float(result[4])],
            "[x3][y3]": [float(result[5]), float(result[6])],
            "[x4][y4]": [float(result[7]), float(result[8])],
        }
        # for i in range(8):
        #     bbox.append(float(result[i+1]))
        # area = (float(result[3]) - float(result[1])) * (float(result[6]) - float(result[4])) * 3.97 * 3.97/10000
        area=calculate_real_area(float(result[1]), float(result[2]),float(result[3]), float(result[4]),float(result[5]), float(result[6]),float(result[7]), float(result[8]),pixel_to_cm)
        data['correct_results'].append({
            "image_name": image_name,
            "datas": {
                 "class_id": int(result[0]),
                 "class_name": result_dict[int(result[0])],
                 "bbox": bbox,
                 "area":area
            }
         })

    return data

# predict_img(r'D:\Work\PythonProjects\wyr\nty\yolo\datas\bus.jpeg')

@cross_origin(supports_credentials=True)
@app.route('/predict', methods=['POST'])
def hello_world():
    try:
        # base_dir = 'datas'
        base_predicted_dir = 'runs'
        try:
            for file in os.listdir(os.path.join('runs', 'obb')):
                shutil.rmtree(os.path.join('runs', 'obb', file))  # 递归删除文件夹
        except:
            pass
        file_path = request.json.get("inputPath")
        out_path = request.json.get("outputPath")
        video_out_path = request.json.get("videoPath")
        model_type = request.json.get("model_type")



        out_path_dir_name = os.path.dirname(out_path)
        if not os.path.exists(out_path_dir_name):
            os.makedirs(out_path_dir_name)
        # files = os.listdir(os.path.join(base_dir, file_path))
        files = os.listdir(os.path.join(file_path))
        results = []
        for file in files:
            # predict_img(r'D:\Work\PythonProjects\wyr\nty\yolo\datas\bus.jpeg')
            # img_path = os.path.join(base_dir, file_path, file)
            img_path = os.path.join(file_path, file)
            result = predict_img(img_path,model_type)
            results.append(result)

        first_file = files[0]
        # first_file_path = os.path.join(base_dir, file_path, first_file)
        first_file_path = os.path.join( file_path, first_file)

        img = cv2.imread(first_file_path)
        (h, w, c) = img.shape
        fps = 1
        out_path_dir_name1 = os.path.dirname(video_out_path)
        if not os.path.exists(out_path_dir_name1):
            os.makedirs(out_path_dir_name1)
        # out = cv2.VideoWriter(video_out_path, cv2.VideoWriter_fourcc(*'XVID'), fps, (w, h), True)
        out = cv2.VideoWriter(video_out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h), True)
        predicted_dirs = os.listdir(os.path.join(base_predicted_dir, 'obb'))
        predicted_images_path = []
        for predicted_dir in predicted_dirs:
            files = os.listdir(os.path.join(base_predicted_dir, 'obb', predicted_dir))
            for file in files:
                if not os.path.isdir(os.path.join(base_predicted_dir, 'obb', predicted_dir, file)):
                    predicted_images_path.append(os.path.join(base_predicted_dir, 'obb', predicted_dir, file))
        for file in predicted_images_path:
            img = cv2.imread(file)
            out.write(img)
        out.release()
        convert_video_to_h264(video_out_path)
        res = {
            'results': results
        }
        with open(out_path, 'w') as f:
            f.write(json.dumps(res, separators=(',', ':')))
        return jsonify({"code": 200, "message": "识别检测成功", "datas": ""})
    except Exception as e:
        msg = str(e)
        return jsonify({"code": 500,  "message": msg, "datas": ""})


def calculate_real_area(x1, y1, x2, y2, x3, y3, x4, y4, pixel_to_cm_ratio):
    """
    Calculate the real area of an oriented bounding box (OBB) given by four vertices,
    where each pixel represents a certain number of centimeters.

    Args:
    coords (list of tuples): List of four tuples, each containing the x and y coordinates of a vertex in pixels.
    pixel_to_cm_ratio (float): Number of centimeters each pixel represents.

    Returns:
    float: The real area of the oriented bounding box in square meters.
    """
    # Calculate the pixel area using the shoelace formula

    pixel_area = abs(
        (x1 * y2 + x2 * y3 + x3 * y4 + x4 * y1) -
        (y1 * x2 + y2 * x3 + y3 * x4 + y4 * x1)
    ) / 2

    # Convert pixel area to square centimeters
    cm_area = pixel_area * (pixel_to_cm_ratio ** 2)

    # Convert square centimeters to square meters
    real_area = cm_area / 10000

    return real_area


# Example usage
# example_coords = [(1, 1), (4, 5), (8, 3), (5, -2)]
# pixel_to_cm = 3.9  # each pixel represents 3.9 cm
# real_area = calculate_real_area(example_coords, pixel_to_cm)
# print(f"The real area of the oriented bounding box is: {real_area:.4f} square meters")



def load_json(file_path):
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")

    with open(file_path, 'r') as file:
        try:
            return json.load(file)
        except json.JSONDecodeError:
            raise ValueError(f"Invalid JSON format: {file_path}")

def extract_bboxes(data):
    if 'results' not in data:
        raise ValueError("Invalid data format: 'results' key not found")

    bboxes = {}
    for item in data['results']:
        if 'correct_results' not in item:
            continue

        for result in item['correct_results']:
            if 'image_name' not in result or 'datas' not in result:
                continue

            image_name = result['image_name']
            suffix = '_'.join(image_name.split('_')[-2:])
            data = result['datas']

            if 'class_id' not in data or 'bbox' not in data or 'area' not in data:
                continue

            class_id = data['class_id']
            class_name = data.get('class_name', '')
            bbox = data['bbox']
            area = data['area']

            if suffix not in bboxes:
                bboxes[suffix] = {}
            if class_id not in bboxes[suffix]:
                bboxes[suffix][class_id] = {'bboxes': [], 'class_name': class_name}
            bboxes[suffix][class_id]['bboxes'].append((image_name, bbox, area))

    return bboxes

def calculate_center(bbox):
    x_center = np.mean([bbox['[x1][y1]'][0], bbox['[x2][y2]'][0], bbox['[x3][y3]'][0], bbox['[x4][y4]'][0]])
    y_center = np.mean([bbox['[x1][y1]'][1], bbox['[x2][y2]'][1], bbox['[x3][y3]'][1], bbox['[x4][y4]'][1]])
    return (x_center, y_center)

def calculate_distance_matrix(centers1, centers2):
    if len(centers1) == 0 or len(centers2) == 0:
        return np.array([])  # 返回空数组以避免错误
    centers1 = np.array(centers1)
    centers2 = np.array(centers2)
    return np.linalg.norm(centers2[:, np.newaxis] - centers1, axis=2)




def compare_bboxes(bboxes1, bboxes2, threshold_count=10):
    results = {}

    common_suffixes = set(bboxes1.keys()).intersection(bboxes2.keys())

    for suffix in common_suffixes:
        image_name1 = bboxes1[suffix][list(bboxes1[suffix].keys())[0]]['bboxes'][0][0]
        image_name2 = bboxes2[suffix][list(bboxes2[suffix].keys())[0]]['bboxes'][0][0]

        added_features = []
        removed_features = []

        for class_id in bboxes2[suffix]:
            bbox_list1 = bboxes1[suffix].get(class_id, {}).get('bboxes', [])
            bbox_list2 = bboxes2[suffix][class_id]['bboxes']
            class_name = bboxes2[suffix][class_id]['class_name']

            centers1 = [calculate_center(b[1]) for b in bbox_list1]
            centers2 = [calculate_center(b[1]) for b in bbox_list2]

            # 计算距离矩阵
            distances = calculate_distance_matrix(centers1, centers2)

            if distances.size == 0:
                if len(centers1) > len(centers2):
                    removed_bboxes = bbox_list1
                    for image_name, bbox, area in removed_bboxes:
                        removed_features.append({
                            'area': area,
                            'bbox': bbox,
                            'class_id': class_id,
                            'class_name': class_name
                        })
                elif len(centers2) > len(centers1):
                    added_bboxes = bbox_list2
                    for image_name, bbox, area in added_bboxes:
                        added_features.append({
                            'area': area,
                            'bbox': bbox,
                            'class_id': class_id,
                            'class_name': class_name
                        })
            else:
                # 找到每个bbox_list2中心点到bbox_list1所有中心点的最小距离
                min_distances = np.min(distances, axis=1)

                # 新增边界框
                if len(centers2) > len(centers1):
                    added_indices = np.argsort(min_distances)[-abs(len(centers2) - len(centers1)):]
                    added_bboxes = [bbox_list2[i] for i in added_indices]
                    for image_name, bbox, area in added_bboxes:
                        added_features.append({
                            'area': area,
                            'bbox': bbox,
                            'class_id': class_id,
                            'class_name': class_name
                        })

                # 减少边界框
                if len(centers1) > len(centers2):
                    distances = calculate_distance_matrix(centers2, centers1)
                    min_distances = np.min(distances, axis=1)
                    removed_indices = np.argsort(min_distances)[-abs(len(centers1) - len(centers2)):]
                    removed_bboxes = [bbox_list1[i] for i in removed_indices]
                    for image_name, bbox, area in removed_bboxes:
                        removed_features.append({
                            'area': area,
                            'bbox': bbox,
                            'class_id': class_id,
                            'class_name': class_name
                        })

        # 数量检查并处理列表
        if len(added_features) < threshold_count:
            added_features = []
        elif len(added_features) > threshold_count:
            # 移除前 threshold_count 个元素，保留剩余部分
            added_features = added_features[threshold_count:]

        if len(removed_features) < threshold_count:
            removed_features = []
        elif len(removed_features) > threshold_count:
            # 移除前 threshold_count 个元素，保留剩余部分
            removed_features = removed_features[threshold_count:]

        # 构建结果格式
        results[suffix] = {
            'image_name1': image_name1,
            'image_name2': image_name2,
            'added_features': added_features,
            'removed_features': removed_features
        }

    fake_heavy_task()

    return list(results.values())


@cross_origin(supports_credentials=True)
@app.route('/compare', methods=['POST'])
def compare():
    try:
        json_path1 = request.json.get("json_path1")
        json_path2 = request.json.get("json_path2")

        if not json_path1 or not json_path2:
            return jsonify({"code": 400, "message": "json_path1 and json_path2 are required"}), 400

        data1 = load_json(json_path1)
        data2 = load_json(json_path2)

        bboxes1 = extract_bboxes(data1)
        bboxes2 = extract_bboxes(data2)

        results = compare_bboxes(bboxes1, bboxes2)

        return jsonify({"code": 200, "data": {"results": results}})

        # json_file_path = r'/home/code/20241211water/ultralytics-main/json/transformed_coordinates.json'  # 替换为你的JSON文件路径
        # if os.path.exists(json_file_path):
        #     with open(json_file_path, 'r', encoding='utf-8') as f:
        #         data = json.load(f)
        #     return jsonify(data)
        # else:
        #     return jsonify({"error": "File not found"}), 404
    except Exception as e:
        # 异常处理：如果出现任何异常，返回 500 错误
        return jsonify({"code": 500, "message": str(e)}), 500


if __name__ == '__main__':
    from gevent import pywsgi

    server = pywsgi.WSGIServer(('0.0.0.0', 5000), app)
    print('服务器已创建，正在监听端口 5000...', flush=True)
    server.serve_forever()


